#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
import matplotlib.pyplot as plt

def mexican_hat(t):
    """Mexican hat wavelet function"""
    return (1 - t**2) * np.exp(-t**2 / 2)

def create_test_signal():
    """Create a simple test signal"""
    t = np.linspace(-5, 5, 200)
    signal = np.zeros_like(t)
    # Create two pulses
    signal += 2.0 * np.exp(-5 * (t + 1.5)**2)
    signal += 1.5 * np.exp(-2 * (t - 2.0)**2)
    return t, signal

def visualize_convolution():
    """Create visualizations of convolution process"""
    # Create signal
    t, signal = create_test_signal()
    dt = t[1] - t[0]
    
    # Create wavelet
    t_wavelet = np.linspace(-3, 3, 100)
    wavelet = mexican_hat(t_wavelet)
    
    # Calculate convolution
    conv_result = np.convolve(signal, wavelet, mode='same') * dt
    
    # Figure 1: Original signal and wavelet
    plt.figure(figsize=(10, 6))
    plt.subplot(211)
    plt.plot(t, signal, 'b-')
    plt.title('Original Signal')
    plt.grid(True)
    
    plt.subplot(212)
    plt.plot(t_wavelet, wavelet, 'r-')
    plt.title('Mexican Hat Wavelet')
    plt.grid(True)
    
    plt.tight_layout()
    plt.savefig('signal_wavelet.png')
    plt.close()
    
    # Figure 2: Convolution result
    plt.figure(figsize=(10, 4))
    plt.plot(t, conv_result, 'g-')
    plt.title('Convolution Result')
    plt.grid(True)
    plt.savefig('convolution_result.png')
    plt.close()
    
    # Figure 3: Key points in convolution process
    plt.figure(figsize=(12, 10))
    
    # Reverse the wavelet for convolution
    reversed_wavelet = np.flip(wavelet)
    m = len(wavelet)
    n = len(signal)
    
    # Select positions to show
    positions = [50, 100, 150]
    
    for i, pos in enumerate(positions):
        plt.subplot(len(positions), 1, i+1)
        
        # Plot signal
        plt.plot(t, signal, 'b-', label='Signal')
        
        # Calculate corresponding t position
        t_pos = t[pos]
        
        # Create wavelet at current position
        wavelet_position = np.linspace(t_pos - 3, t_pos + 3, 100)
        plt.plot(wavelet_position, reversed_wavelet, 'r-', label='Wavelet')
        
        # Calculate overlapping region
        overlap_start = max(0, pos - m // 2)
        overlap_end = min(n, pos + m // 2 + 1)
        wavelet_start = max(0, m // 2 - pos)
        wavelet_end = wavelet_start + (overlap_end - overlap_start)
        
        # Calculate product only in overlapping region
        product = np.zeros_like(signal)
        product[overlap_start:overlap_end] = signal[overlap_start:overlap_end] * reversed_wavelet[wavelet_start:wavelet_end]
        
        # Plot product
        plt.fill_between(t, product, color='g', alpha=0.3, label='Product')
        
        # Calculate convolution value
        conv_value = np.sum(product) * dt
        plt.title(f'Convolution at t={t_pos:.2f}, Value={conv_value:.4f}')
        plt.legend()
        plt.grid(True)
    
    plt.tight_layout()
    plt.savefig('convolution_steps.png')
    plt.close()
    
    print("\nGenerated visualization files:")
    print("1. signal_wavelet.png")
    print("2. convolution_result.png")
    print("3. convolution_steps.png")

def explain_convolution_process():
    """Explain the convolution process"""
    print("\n===== CONVOLUTION CALCULATION PROCESS =====")
    print("\n1. PRINCIPLE:")
    print("   Convolution combines two functions by sliding one over the other")
    print("   and calculating the area under their product at each position.")
    print("\n2. STEPS:")
    print("   a. Reverse the wavelet function (time reversal)")
    print("   b. Slide the reversed wavelet across the signal")
    print("   c. At each position, multiply corresponding values")
    print("   d. Sum all products to get the convolution value")
    print("\n3. MEXICAN HAT WAVELET:")
    print("   - Formula: (1 - t^2) * exp(-t^2 / 2)")
    print("   - Second derivative of Gaussian")
    print("   - Good for edge detection and feature extraction")
    print("\n4. SAMPLE CALCULATION:")
    print("   Example for a few points:")
    print("   Suppose signal = [s1, s2, s3, s4], wavelet = [w1, w2, w3]")
    print("   reversed wavelet = [w3, w2, w1]")
    print("   conv[0] = s1*w1")
    print("   conv[1] = s1*w2 + s2*w1")
    print("   conv[2] = s1*w3 + s2*w2 + s3*w1")
    print("   conv[3] = s2*w3 + s3*w2 + s4*w1")
    print("   conv[4] = s3*w3 + s4*w2")
    print("   conv[5] = s4*w3")

def main():
    print("Running Mexican hat wavelet convolution demonstration...")
    visualize_convolution()
    explain_convolution_process()
    print("\nDemonstration completed successfully!")

if __name__ == "__main__":
    main()